package com.treasure.web.interceptor;

import com.google.common.cache.Cache;
import com.google.common.cache.CacheBuilder;
import com.treasure.common.exception.RateLimitException;
import org.jetbrains.annotations.NotNull;
import org.springframework.web.servlet.HandlerInterceptor;
import org.springframework.web.servlet.ModelAndView;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.util.Objects;
import java.util.concurrent.TimeUnit;

public class RateLimitInterceptor implements HandlerInterceptor {

    private static final int MAX_REQUESTS_PER_IP = 2; // 允许的最大请求次数
    private static final int CACHE_EXPIRATION_MINUTES = 5; // 缓存过期时间

    private final Cache<String, Integer> requestCountCache;

    public RateLimitInterceptor() {
        requestCountCache = CacheBuilder.newBuilder()
                .expireAfterWrite(CACHE_EXPIRATION_MINUTES, TimeUnit.SECONDS)
                .build();
    }

    @Override
    public boolean preHandle(@NotNull HttpServletRequest request, @NotNull HttpServletResponse response, @NotNull Object handler) throws Exception {
        String clientIpAddress = getClientIP(request);
        String url = request.getRequestURL().toString();
        String limitCacheKey = "["+clientIpAddress+"]-["+url+"]";
        Integer requestCount = requestCountCache.getIfPresent(limitCacheKey);
        if (Objects.isNull(requestCount)) {
            requestCount = 0;
        }

        if (requestCount >= MAX_REQUESTS_PER_IP) {
            response.setStatus(429);
            response.setCharacterEncoding("UTF-8"); // 设置字符编码为UTF-8
            response.setContentType("application/json;charset=UTF-8"); // 设置Content-Type及字符编码
            throw new RateLimitException("请求过于频繁，请稍后再试。");
        }

        requestCountCache.put(limitCacheKey, requestCount + 1);
        return true;
    }

    @Override
    public void postHandle(HttpServletRequest request, HttpServletResponse response, Object handler, ModelAndView modelAndView) throws Exception {
        // 这里可以添加后处理逻辑
    }

    @Override
    public void afterCompletion(@NotNull HttpServletRequest request, @NotNull HttpServletResponse response, @NotNull Object handler, Exception ex) throws Exception {
        // 这里可以添加完成处理逻辑
    }

    private String getClientIP(HttpServletRequest request) {
        String ipAddress = request.getHeader("X-Forwarded-For");
        if (ipAddress == null || ipAddress.isEmpty() || "unknown".equalsIgnoreCase(ipAddress)) {
            ipAddress = request.getHeader("Proxy-Client-IP");
        }
        if (ipAddress == null || ipAddress.isEmpty() || "unknown".equalsIgnoreCase(ipAddress)) {
            ipAddress = request.getHeader("WL-Proxy-Client-IP");
        }
        if (ipAddress == null || ipAddress.isEmpty() || "unknown".equalsIgnoreCase(ipAddress)) {
            ipAddress = request.getRemoteAddr();
        }
        return ipAddress;
    }
}
